#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/frame.h>
#include <libavutil/pixfmt.h>
#include <libavcodec/codec.h>
#include <libavcodec/packet.h>
#include <libavutil/avutil.h>
#include <libswscale/swscale.h>
#include <libswresample/swresample.h>
#include <libavutil/mem.h>
#include <libavutil/samplefmt.h>
#include <stdint.h>
#include <stdio.h>
#include <sys/types.h>
#include <libavutil/pixfmt.h>
#include <libavutil/imgutils.h>

static int get_format_from_sample_fmt(const char **fmt, enum AVSampleFormat sample_fmt) {
    struct sample_fmt_entry {
        enum AVSampleFormat sample_fmt; const char *fmt_be, *fmt_le;
    } sample_fmt_entries[] = {
        { AV_SAMPLE_FMT_U8,  "u8",    "u8"    },
        { AV_SAMPLE_FMT_S16, "s16be", "s16le" },
        { AV_SAMPLE_FMT_S32, "s32be", "s32le" },
        { AV_SAMPLE_FMT_FLT, "f32be", "f32le" },
        { AV_SAMPLE_FMT_DBL, "f64be", "f64le" },
    };
    *fmt = NULL;

    for (int i = 0; i < FF_ARRAY_ELEMS(sample_fmt_entries); i++) {
        struct sample_fmt_entry *entry = &sample_fmt_entries[i];
        if (sample_fmt == entry->sample_fmt) {
            *fmt = AV_NE(entry->fmt_be, entry->fmt_le);
            return 0;
        }
    }

    fprintf(stderr,
            "sample format %s is not supported as output format\n",
            av_get_sample_fmt_name(sample_fmt));
    return -1;
}

static void print_pcm_info(AVCodecContext* dec_ctx, const char* output) {
    const char* fmt;
    enum AVSampleFormat sfmt = dec_ctx->sample_fmt;
    if (av_sample_fmt_is_planar(sfmt)) {
        sfmt = av_get_packed_sample_fmt(sfmt);
    }
    if (get_format_from_sample_fmt(&fmt, sfmt) < 0) {
        return;
    }
    printf("Play the output audio file with the command:\n\t"
        "ffplay -f %s -ac %d -ar %d %s\n",
        fmt, dec_ctx->ch_layout.nb_channels, dec_ctx->sample_rate, output);
}

void print_yuv_info(AVCodecContext* dec_ctx, const char* output) {
    printf("Play the output video file with the command:\n\t"
        "ffplay -video_size %dx%d -pixel_format %s %s\n",
        dec_ctx->width, dec_ctx->height,
        av_get_pix_fmt_name(dec_ctx->pix_fmt), output);
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <input file>\n", argv[0]);
        exit(0);
    }
    
    const char* input = argv[1];
    const char* output_yuv = "output.yuv";
    const char* output_pcm = "output.pcm";

    // Open input file
    FILE* fin = fopen(input, "rb");
    if (!fin) {
        fprintf(stderr, "Failed to open input file %s\n", input);
        exit(1);
    }
    // Open output file
    FILE* fout_pcm = fopen(output_pcm, "wb");
    if (!fout_pcm) {
        fprintf(stderr, "Failed to open output file %s\n", output_pcm);
        exit(1);
    }
    FILE* fout_yuv = fopen(output_yuv, "wb");
    if (!fout_yuv) {
        fprintf(stderr, "Failed to open output file %s\n", output_yuv);
        exit(1);
    }

    // Open an input stream and read the header
    AVFormatContext* fmt_ctx = NULL;
    if (avformat_open_input(&fmt_ctx, input, NULL, NULL) < 0) {
        fprintf(stderr, "Failed to open input stream\n");
        exit(1);
    }
    
    // Retrieve stream information
    if (avformat_find_stream_info(fmt_ctx, NULL) < 0) {
        fprintf(stderr, "Failed to find stream info\n");
        exit(1);
    }
    
    // Find video and audio stream index
    int vs_idx = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0);
    int as_idx = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_AUDIO, -1, -1, NULL, 0);
    if (vs_idx < 0 || as_idx < 0) {
        fprintf(stderr, "Failed to find video or audio stream\n");
        return -1;
    }
    
    // Get video and audio stream
    AVStream* vs = fmt_ctx->streams[vs_idx];
    AVStream* as = fmt_ctx->streams[as_idx];

    // Find video and audio decoder
    const AVCodec* vdec = avcodec_find_decoder(vs->codecpar->codec_id);
    if (!vdec) {
        fprintf(stderr, "Failed to find video decoder\n");
        return -1;
    }
    const AVCodec* adec = avcodec_find_decoder(as->codecpar->codec_id);
    if (!adec) {
        fprintf(stderr, "Failed to find audio decoder\n");
        return -1;
    }

    // Allocate video and audio codec context
    AVCodecContext* vdec_ctx = avcodec_alloc_context3(vdec);
    if (!vdec_ctx) {
        fprintf(stderr, "Failed to allocate video codec context\n");
        return -1;
    }
    AVCodecContext* adec_ctx = avcodec_alloc_context3(adec);
    if (!adec_ctx) {
        fprintf(stderr, "Failed to allocate audio codec context\n");
        return -1;
    }

    // Copy codec parameters from input stream to codec context
    if (avcodec_parameters_to_context(vdec_ctx, vs->codecpar) < 0) {
        fprintf(stderr, "Failed to copy video codec parameters to codec context\n");
        return -1;
    }
    if (avcodec_parameters_to_context(adec_ctx, as->codecpar) < 0) {
        fprintf(stderr, "Failed to copy audio codec parameters to codec context\n");
        return -1;
    }

    // Initialize video and audio codec context to use the given codec
    if (avcodec_open2(vdec_ctx, vdec, NULL) < 0) {
        fprintf(stderr, "Failed to initialize video codec context\n");
        return -1;
    }
    if (avcodec_open2(adec_ctx, adec, NULL) < 0) {
        fprintf(stderr, "Failed to initialize audio codec context\n");
        return -1;
    }

    // Allocate an packet and frame
    AVPacket* pkt = av_packet_alloc();
    if (!pkt) {
        fprintf(stderr, "Failed to allocate audio packet\n");
        exit(1);
    }
    AVFrame* frame = av_frame_alloc();
    if (!frame) {
        fprintf(stderr, "Failed to allocate audio frame\n");
        exit(1);
    }
    
    int ret = 0;
    int is_audio = 0;
    AVCodecContext* dec_ctx = NULL;
    while (1) {
        // Read packet from input file
        ret = av_read_frame(fmt_ctx, pkt);
        if (ret < 0) break;

        printf("=============================\n");
        printf("Packet size: %d\n", pkt->size);
        printf("Packet pts: %ld\n", pkt->pts);
        printf("Packet dts: %ld\n", pkt->dts);
        printf("Packet duration: %ld\n", pkt->duration);

        // Select codec context
        dec_ctx = NULL;
        if (pkt->stream_index == vs_idx) {
            printf("Decoding video packet\n");
            dec_ctx = vdec_ctx;
        }
        if (pkt->stream_index == as_idx) {
            printf("Decoding audio packet\n");
            dec_ctx = adec_ctx;
        }
        if (!dec_ctx) {
            fprintf(stderr, "Failed to parse packet\n");
            continue;
        }

        // Submit packet to decoder
        ret = avcodec_send_packet(dec_ctx, pkt);
        av_packet_unref(pkt);
        if (ret != 0) {
            fprintf(stderr, "Error submitting the packet to the decoder\n");
            exit(1);
        }

        // Receive frame from decoder
        while (ret >= 0) {
            ret = avcodec_receive_frame(dec_ctx, frame);
            if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
                break;
            }
            if (ret < 0) {
                fprintf(stderr, "Error during decoding\n");
                exit(1);
            }
            
            // Write PCM data to output file
            if (dec_ctx == adec_ctx) {
                int data_size = av_get_bytes_per_sample(adec_ctx->sample_fmt);
                if (data_size < 0) {
                    fprintf(stderr, "Failed to calculate data size\n");
                    exit(1);
                }
                for (int i = 0; i < frame->nb_samples; i++) {
                    for (int ch = 0; ch < adec_ctx->ch_layout.nb_channels; ch++) {
                        fwrite(frame->data[ch] + data_size*i, 1, data_size, fout_pcm);
                    }
                }
            }
            // Write YUV data to output file
            else {
                // Y: frame->data[0]
                for (int i = 0; i < frame->height; i++) {
                    fwrite(frame->data[0] + i * frame->linesize[0], 1, frame->width, fout_yuv);
                }
                // U: frame->data[1]
                for (int i = 0; i < frame->height / 2; i++) {
                    fwrite(frame->data[1] + i * frame->linesize[1], 1, frame->width / 2, fout_yuv);
                }
                // V: frame->data[2]
                for (int i = 0; i < frame->height / 2; i++) {
                    fwrite(frame->data[2] + i * frame->linesize[2], 1, frame->width / 2, fout_yuv);
                }
            }
        }
        printf("=============================\n");
    }

    // Print pcm information
    print_pcm_info(adec_ctx, output_pcm);
    // Print yuv information
    print_yuv_info(vdec_ctx, output_yuv);

    // Clean
    av_frame_free(&frame);
    av_packet_free(&pkt);
    avcodec_free_context(&adec_ctx);
    avformat_close_input(&fmt_ctx);
    fclose(fin);
    fclose(fout_pcm);

    return 0;
}
